# trainer/training_controller.py
import torch
import numpy as np
from typing import Dict, List, Optional
import asyncio
from datetime import datetime
import json
import os
import sys
import torch.nn.functional as F
from enum import Enum

# 添加项目根目录到Python路径
from pathlib import Path
sys.path.append('/app')
try:
    from models.lightweight_ppo import MultiHeadPPO
    from game_engine.game_state import FactorType
    from models.model_strategy_manager import ModelStrategyManager, create_default_strategy_manager
    from utils.context_manager import ContextManager, ContextPriority
except ImportError as e:
    print(f"导入错误: {e}")
    # 备用导入路径
    sys.path.append(str(Path(__file__).parent.parent))  # 添加项目根目录
    from models.lightweight_ppo import MultiHeadPPO
    from game_engine.game_state import FactorType
    from models.model_strategy_manager import ModelStrategyManager, create_default_strategy_manager
    from utils.context_manager import ContextManager, ContextPriority
from training.experience_replay import PrioritizedExperienceReplay
from training.curriculum_learning import CurriculumManager
from utils.training_monitor import TrainingMonitor

class TrainingController:
    """训练控制器 - 管理整个训练流程"""
    
    def __init__(self, config: Dict, online_api_key: Optional[str] = None):
        self.config = config
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # 初始化组件
        self.ppo = MultiHeadPPO(
            factor_types=config['factor_types'],
            lr=config.get('learning_rate', 1e-4),
            gamma=config.get('gamma', 0.99),
            clip_epsilon=config.get('clip_epsilon', 0.2)
        )
        
        self.experience_replay = PrioritizedExperienceReplay(
            capacity=config.get('replay_capacity', 10000),
            alpha=config.get('priority_alpha', 0.6),
            beta=config.get('priority_beta', 0.4)
        )
        
        self.curriculum = CurriculumManager(config.get('curriculum_settings', {}))
        self.monitor = TrainingMonitor()
        
        # 初始化策略管理器
        factor_types = config['factor_types']
        self.strategy_manager = create_default_strategy_manager(factor_types, online_api_key)
        
        # 初始化上下文管理器
        self.context_manager = ContextManager()
        
        # 训练状态
        self.current_cycle = 0
        self.best_performance = 0.0
        self.training_history = []
        
        # 创建保存目录
        os.makedirs('models/checkpoints', exist_ok=True)
        os.makedirs('training_logs', exist_ok=True)
    
    def _get_world_state(self, game_world) -> Dict:
        """获取游戏世界状态"""
        return {
            'iron_curtain_progress': game_world.iron_curtain_progress,
            'lygus_health': game_world.lygus.health,
            'current_time': game_world.current_time,
            'total_time': game_world.total_time,
            'alive_factors_count': sum(1 for f in game_world.factors if f.is_alive),
            'game_status': 'active' if game_world.iron_curtain_progress < 100 and game_world.lygus.health > 0 else 'ended'
        }

    def _convert_decision_to_actions(self, decision_dict: Dict) -> Dict:
        """将策略管理器的决策转换为动作格式"""
        # 如果决策已经是动作格式，直接返回
        if 'movement' in decision_dict and 'combat' in decision_dict:
            return {k: torch.tensor(v) if isinstance(v, (int, float)) else torch.tensor(0) for k, v in decision_dict.items()}
        
        # 否则创建默认动作
        return {
            'movement': torch.tensor(decision_dict.get('movement', 0)),
            'combat': torch.tensor(decision_dict.get('combat', 0)),
            'support': torch.tensor(decision_dict.get('support', 0)),
            'special': torch.tensor(decision_dict.get('special', 0))
        }

    def _execute_actions_with_decision(self, factor, decision_result: Dict, game_world) -> float:
        """根据决策结果执行动作并计算奖励"""
        action_dict = decision_result.get("action", {})
        actions = self._convert_decision_to_actions(action_dict)
        return self._execute_actions(factor, actions, game_world)

    def _execute_default_action(self, factor, game_world) -> float:
        """执行默认动作（当决策失败时）"""
        # 创建默认动作（不移动，不攻击，不支援，不使用特殊能力）
        default_actions = {
            'movement': torch.tensor(0),
            'combat': torch.tensor(0),
            'support': torch.tensor(0),
            'special': torch.tensor(0)
        }
        return self._execute_actions(factor, default_actions, game_world)
    
    async def run_training_cycle(self, game_world, cycle_number: int) -> Dict:
        """运行一个训练周期"""
        self.current_cycle = cycle_number
        print(f"开始第 {cycle_number} 轮回训练...")
        
        # 更新课程难度
        current_difficulty = self.curriculum.get_difficulty(cycle_number)
        self._apply_difficulty_settings(game_world, current_difficulty)
        
        cycle_results = {
            'cycle_number': cycle_number,
            'start_time': datetime.now(),
            'difficulty': current_difficulty,
            'episode_results': [],
            'training_metrics': {}
        }
        
        # 运行多个训练回合
        for episode in range(self.config.get('episodes_per_cycle', 10)):
            episode_result = await self._run_training_episode(game_world, episode)
            cycle_results['episode_results'].append(episode_result)
            
            # 每3个回合更新一次策略
            if episode % 3 == 0 and len(self.experience_replay) > self.config.get('min_batch_size', 32):
                training_metrics = self._update_policies()
                cycle_results['training_metrics'].update(training_metrics)
        
        # 周期结束处理
        final_metrics = self._evaluate_cycle_performance(cycle_results)
        cycle_results.update(final_metrics)
        
        # 保存检查点
        if final_metrics['overall_performance'] > self.best_performance:
            self.best_performance = final_metrics['overall_performance']
            self._save_checkpoint(cycle_number, final_metrics['overall_performance'])
        
        # 记录训练历史
        self.training_history.append(cycle_results)
        self._save_training_log(cycle_results)
        
        print(f"第 {cycle_number} 轮回训练完成. 性能: {final_metrics['overall_performance']:.3f}")
        return cycle_results
    
    async def _run_training_episode(self, game_world, episode_num: int) -> Dict:
        """运行单个训练回合"""
        episode_experiences = []
        total_reward = 0
        steps = 0
        
        # 重置游戏世界（保持智能体策略）
        self._reset_episode(game_world)
        
        # 运行回合
        max_steps = self.config.get('max_steps_per_episode', 200)
        
        for step in range(max_steps):
            step_experiences, step_reward = await self._run_training_step(game_world, step)
            episode_experiences.extend(step_experiences)
            total_reward += step_reward
            steps += 1
            
            # 检查回合结束条件
            if self._should_end_episode(game_world):
                break
        
        # 计算回合回报
        discounted_returns = self._compute_discounted_returns(episode_experiences)
        
        # 添加经验到回放缓冲区
        for exp, ret in zip(episode_experiences, discounted_returns):
            exp['returns'] = ret
            # 计算优势函数
            exp['advantages'] = ret - exp['value'].item()
            self.experience_replay.add(exp, abs(exp['advantages']))
        
        return {
            'episode': episode_num,
            'total_reward': total_reward,
            'steps': steps,
            'surviving_factors': sum(1 for f in game_world.factors if f.is_alive),
            'lygus_health': game_world.lygus.health,
            'iron_curtain_progress': game_world.iron_curtain_progress
        }
    
    async def _run_training_step(self, game_world, step_num: int) -> tuple:
        """运行单个训练步骤 - 更新以使用策略管理器"""
        step_experiences = []
        step_reward = 0
        
        # 更新游戏世界
        game_world.update_world(1.0)
        
        # 为每个因子生成经验
        for factor in game_world.factors:
            if not factor.is_alive:
                continue
            
            # 获取当前状态
            state = self._get_factor_state(factor, game_world)
            
            # 使用策略管理器获取决策
            decision_result = await self.strategy_manager.execute_decision(
                component_name=factor.factor_type.name.lower(),
                state=state,
                decision_context={
                    "current_time": step_num,
                    "game_world_state": self._get_world_state(game_world),
                    "factor_status": factor.to_dict() if hasattr(factor, 'to_dict') else {
                        'name': factor.name,
                        'level': factor.level,
                        'health': factor.attributes['health'],
                        'location': factor.current_location,
                        'is_alive': factor.is_alive
                    }
                }
            )
            
            # 执行决策
            if decision_result["success"]:
                # 将策略管理器的决策转换为动作格式
                action_dict = decision_result.get("action", {})
                actions = self._convert_decision_to_actions(action_dict)
                
                # 计算奖励
                reward = self._execute_actions_with_decision(factor, decision_result, game_world)
                step_reward += reward
                
                # 存储经验
                experience = {
                    'factor_type': factor.factor_type,
                    'state': torch.FloatTensor(state).to(self.device),
                    'actions': actions,
                    'old_log_probs': decision_result.get('log_probs', {}),
                    'value': decision_result.get('value', 0),
                    'reward': reward,
                    'step': step_num,
                    'model_used': decision_result.get('model_used', 'unknown')
                }
                step_experiences.append(experience)
            else:
                # 处理决策失败的情况
                reward = self._execute_default_action(factor, game_world)
                step_reward += reward
        
        return step_experiences, step_reward
    
    def _get_factor_state(self, factor, game_world) -> np.ndarray:
        """获取因子状态表示"""
        state = []
        
        # 自身属性
        state.extend([
            factor.level / 20.0,  # 标准化等级
            factor.attributes['attack'] / 100.0,
            factor.attributes['defense'] / 100.0,
            factor.attributes['health'] / 200.0,
            factor.attributes['healing'] / 100.0,
            factor.attributes['mobility'] / 50.0,
            factor.attributes['perception'] / 50.0
        ])
        
        # 位置信息
        state.extend([
            factor.current_location[0] / 100.0,
            factor.current_location[1] / 100.0
        ])
        
        # 附近环境信息（简化）
        nearby_factors = self._get_nearby_factors(factor, game_world, radius=20.0)
        state.append(len(nearby_factors) / 10.0)  # 附近因子数量
        
        # 游戏状态
        state.extend([
            game_world.iron_curtain_progress / 100.0,
            game_world.lygus.health / 10000.0,
            game_world.current_time / game_world.total_time
        ])
        
        # 填充到固定维度
        while len(state) < 128:
            state.append(0.0)
        
        return np.array(state[:128], dtype=np.float32)
    
    def _get_nearby_factors(self, factor, game_world, radius: float):
        """获取附近的因子"""
        nearby = []
        for other in game_world.factors:
            if other != factor and other.is_alive:
                distance = np.sqrt(
                    (factor.current_location[0] - other.current_location[0])**2 +
                    (factor.current_location[1] - other.current_location[1])**2
                )
                if distance <= radius:
                    nearby.append(other)
        return nearby
    
    def _execute_actions(self, factor, actions: Dict, game_world) -> float:
        """执行动作并计算奖励"""
        reward = 0.0
        
        # 解码动作并执行
        # 移动动作
        move_action = actions['movement'].item()
        new_location = self._decode_movement(factor.current_location, move_action)
        factor.move_to(new_location)
        
        # 战斗动作
        combat_action = actions['combat'].item()
        reward += self._execute_combat_action(factor, combat_action, game_world)
        
        # 支援动作
        support_action = actions['support'].item()
        reward += self._execute_support_action(factor, support_action, game_world)
        
        # 特殊动作
        special_action = actions['special'].item()
        reward += self._execute_special_action(factor, special_action, game_world)
        
        return reward
    
    def _decode_movement(self, current_location, action: int) -> tuple:
        """解码移动动作"""
        directions = [
            (0, 1), (1, 1), (1, 0), (1, -1),
            (0, -1), (-1, -1), (-1, 0), (-1, 1)
        ]
        
        if action < len(directions):
            dx, dy = directions[action]
            new_x = max(0, min(100, current_location[0] + dx * 10))
            new_y = max(0, min(100, current_location[1] + dy * 10))
            return (new_x, new_y)
        
        return current_location
    
    def _execute_combat_action(self, factor, action: int, game_world) -> float:
        """执行战斗动作"""
        reward = 0.0
        
        # 查找附近敌人
        nearby_enemies = [f for f in self._get_nearby_factors(factor, game_world, 15.0) 
                         if self._is_enemy(f, factor)]
        
        if nearby_enemies and action > 0:  # 非"不攻击"动作
            target = nearby_enemies[0]  # 简化：攻击第一个敌人
            damage = factor.attributes['attack'] * (0.5 + action * 0.1)
            actual_damage = target.take_damage(damage)
            reward += actual_damage * 0.1
            
            if not target.is_alive:
                reward += 50.0  # 击败奖励
        
        return reward
    
    def _execute_support_action(self, factor, action: int, game_world) -> float:
        """执行支援动作"""
        reward = 0.0
        
        # 查找附近盟友
        nearby_allies = [f for f in self._get_nearby_factors(factor, game_world, 15.0) 
                        if not self._is_enemy(f, factor) and f != factor]
        
        if nearby_allies and action > 0:
            target = nearby_allies[0]  # 简化：支援第一个盟友
            
            if action == 1:  # 治疗
                heal_amount = factor.attributes['healing'] * 0.8
                target.attributes['health'] = min(
                    target.attributes['health'] + heal_amount,
                    target.attributes['health']  # 最大生命值，这里需要实际的最大值
                )
                reward += heal_amount * 0.05
            
            # 其他支援动作...
        
        return reward
    
    def _execute_special_action(self, factor, action: int, game_world) -> float:
        """执行特殊动作"""
        reward = 0.0
        
        if action > 0 and factor.special_charges > 0:
            factor.use_special_ability()
            reward += 10.0  # 使用特殊能力的基础奖励
        
        return reward
    
    def _is_enemy(self, factor1, factor2) -> bool:
        """判断两个因子是否为敌对关系"""
        # 简化：所有因子互相敌对（在PVP环境中）
        return factor1 != factor2
    
    def _update_policies(self) -> Dict:
        """更新策略网络"""
        if len(self.experience_replay) < self.config.get('min_batch_size', 32):
            return {}
        
        # 采样经验批次
        batch_size = min(self.config.get('batch_size', 64), len(self.experience_replay))
        batch, indices, weights = self.experience_replay.sample(batch_size)
        
        # 更新策略
        update_metrics = self.ppo.update(batch)
        
        # 更新优先级
        new_priorities = [abs(exp['advantages']) for exp in batch]
        self.experience_replay.update_priorities(indices, new_priorities)
        
        return update_metrics
    
    def _compute_discounted_returns(self, experiences: List) -> List[float]:
        """计算折扣回报"""
        returns = []
        current_return = 0
        
        # 反向计算
        for exp in reversed(experiences):
            current_return = exp['reward'] + self.ppo.gamma * current_return
            returns.insert(0, current_return)
        
        # 标准化回报
        returns = np.array(returns)
        if len(returns) > 1:
            returns = (returns - returns.mean()) / (returns.std() + 1e-8)
        
        return returns.tolist()
    
    def _should_end_episode(self, game_world) -> bool:
        """判断是否应该结束回合"""
        # 所有因子死亡
        if all(not factor.is_alive for factor in game_world.factors):
            return True
        
        # 铁幕进度达到100%
        if game_world.iron_curtain_progress >= 100:
            return True
        
        # 来古士被击败
        if game_world.lygus.health <= 0:
            return True
        
        return False
    
    def _reset_episode(self, game_world):
        """重置回合状态"""
        # 重置因子状态
        for factor in game_world.factors:
            factor.is_alive = True
            factor.attributes['health'] = 100  # 重置生命值
            factor.current_location = (np.random.uniform(0, 100), np.random.uniform(0, 100))
            factor.special_charges = 3
        
        # 重置来古士
        game_world.lygus.health = 10000
        game_world.lygus.is_awakened = False
        
        # 重置铁幕进度
        game_world.iron_curtain_progress = 0
        game_world.current_time = 0
    
    def _apply_difficulty_settings(self, game_world, difficulty: Dict):
        """应用难度设置"""
        # 调整来古士属性
        game_world.lygus.attack = 500 * difficulty.get('enemy_strength', 1.0)
        game_world.lygus.defense = 300 * difficulty.get('enemy_strength', 1.0)
        
        # 调整铁幕增长速度
        game_world.total_time = 1000 / difficulty.get('time_pressure', 1.0)
    
    def _evaluate_cycle_performance(self, cycle_results: Dict) -> Dict:
        """评估周期性能"""
        episode_results = cycle_results['episode_results']
        
        if not episode_results:
            return {'overall_performance': 0.0}
        
        # 计算平均奖励
        avg_reward = np.mean([ep['total_reward'] for ep in episode_results])
        
        # 计算生存率
        avg_survival = np.mean([ep['surviving_factors'] for ep in episode_results]) / 12.0
        
        # 计算对来古士造成的伤害
        avg_damage = np.mean([10000 - ep['lygus_health'] for ep in episode_results]) / 10000.0
        
        # 综合性能评分
        performance = (avg_reward * 0.4 + avg_survival * 0.3 + avg_damage * 0.3)
        
        return {
            'overall_performance': performance,
            'average_reward': avg_reward,
            'average_survival_rate': avg_survival,
            'average_damage_rate': avg_damage
        }
    
    def _save_checkpoint(self, cycle: int, performance: float):
        """保存检查点"""
        checkpoint = {
            'cycle': cycle,
            'performance': performance,
            'policies_state': {
                factor_type.name: policy.state_dict()
                for factor_type, policy in self.ppo.policies.items()
            },
            'optimizers_state': {
                factor_type.name: optimizer.state_dict()
                for factor_type, optimizer in self.ppo.optimizers.items()
            },
            'training_config': self.config
        }
        
        filename = f"models/checkpoints/cycle_{cycle}_perf_{performance:.3f}.pth"
        torch.save(checkpoint, filename)
        print(f"检查点已保存: {filename}")
    
    def _save_training_log(self, cycle_results: Dict):
        """保存训练日志"""
        log_entry = {
            'timestamp': datetime.now().isoformat(),
            'cycle': cycle_results['cycle_number'],
            'performance': cycle_results.get('overall_performance', 0),
            'metrics': cycle_results.get('training_metrics', {})
        }
        
        log_file = "training_logs/training_history.json"
        
        # 读取现有日志
        if os.path.exists(log_file):
            with open(log_file, 'r') as f:
                logs = json.load(f)
        else:
            logs = []
        
        # 添加新日志
        logs.append(log_entry)
        
        # 保存日志
        with open(log_file, 'w') as f:
            json.dump(logs, f, indent=2)
